"""WebSocket support for real-time vector operations."""

import asyncio
import json
import logging
import uuid
from typing import Dict, Set, Any
from datetime import datetime

from fastapi import WebSocket, WebSocketDisconnect
from api.models import WebSocketMessage, ErrorResponse
from core.database import get_db_manager
from core.embedding import get_embedding_service, embed_texts, embed_text

logger = logging.getLogger(__name__)


class ConnectionManager:
    """WebSocket connection manager."""

    def __init__(self):
        self.active_connections: Dict[str, WebSocket] = {}
        self.connection_metadata: Dict[str, Dict] = {}

    async def connect(self, websocket: WebSocket, client_id: str = None) -> str:
        """Accept new WebSocket connection."""
        await websocket.accept()
        connection_id = client_id or str(uuid.uuid4())
        self.active_connections[connection_id] = websocket
        self.connection_metadata[connection_id] = {
            "connected_at": datetime.now().isoformat(),
            "message_count": 0
        }
        logger.info(f"WebSocket connected: {connection_id}")
        return connection_id

    def disconnect(self, connection_id: str):
        """Remove WebSocket connection."""
        if connection_id in self.active_connections:
            del self.active_connections[connection_id]
            if connection_id in self.connection_metadata:
                del self.connection_metadata[connection_id]
            logger.info(f"WebSocket disconnected: {connection_id}")

    async def send_message(self, connection_id: str, message: Dict[str, Any]):
        """Send message to specific connection."""
        if connection_id in self.active_connections:
            websocket = self.active_connections[connection_id]
            try:
                await websocket.send_json(message)
                self.connection_metadata[connection_id]["message_count"] += 1
            except Exception as e:
                logger.error(f"Failed to send message to {connection_id}: {e}")
                self.disconnect(connection_id)

    async def broadcast_message(self, message: Dict[str, Any], exclude: Set[str] = None):
        """Broadcast message to all connections."""
        exclude = exclude or set()
        disconnected = []

        for connection_id, websocket in self.active_connections.items():
            if connection_id in exclude:
                continue

            try:
                await websocket.send_json(message)
                self.connection_metadata[connection_id]["message_count"] += 1
            except Exception as e:
                logger.error(f"Failed to broadcast to {connection_id}: {e}")
                disconnected.append(connection_id)

        # Clean up disconnected clients
        for connection_id in disconnected:
            self.disconnect(connection_id)

    def get_connection_stats(self) -> Dict[str, Any]:
        """Get connection statistics."""
        return {
            "total_connections": len(self.active_connections),
            "connections": {
                conn_id: {
                    "connected_at": metadata["connected_at"],
                    "message_count": metadata["message_count"]
                }
                for conn_id, metadata in self.connection_metadata.items()
            }
        }


# Global connection manager
connection_manager = ConnectionManager()


class WebSocketHandler:
    """WebSocket message handler."""

    def __init__(self):
        self.handlers = {
            "embed": self.handle_embed,
            "search": self.handle_search,
            "store": self.handle_store,
            "health": self.handle_health,
            "stats": self.handle_stats,
            "ping": self.handle_ping,
            "subscribe": self.handle_subscribe,
            "unsubscribe": self.handle_unsubscribe
        }
        self.subscriptions: Dict[str, Set[str]] = {}  # connection_id -> set of subscription types

    async def handle_message(self, connection_id: str, message: WebSocketMessage):
        """Handle incoming WebSocket message."""
        try:
            message_type = message.type
            request_id = message.request_id

            if message_type not in self.handlers:
                await self.send_error(connection_id, f"Unknown message type: {message_type}", request_id)
                return

            handler = self.handlers[message_type]
            await handler(connection_id, message.data or {}, request_id)

        except Exception as e:
            logger.error(f"Error handling WebSocket message: {e}")
            await self.send_error(connection_id, f"Internal error: {str(e)}", message.request_id)

    async def handle_embed(self, connection_id: str, data: Dict[str, Any], request_id: str):
        """Handle embedding request."""
        try:
            texts = data.get("texts", [])
            provider = data.get("provider")

            if not texts:
                await self.send_error(connection_id, "No texts provided", request_id)
                return

            # Generate embeddings
            embeddings = await embed_texts(texts, provider)

            response = {
                "type": "embed_response",
                "request_id": request_id,
                "data": {
                    "embeddings": embeddings,
                    "count": len(embeddings),
                    "dimensions": len(embeddings[0]) if embeddings else 0
                },
                "timestamp": datetime.now().isoformat()
            }

            await connection_manager.send_message(connection_id, response)

        except Exception as e:
            logger.error(f"WebSocket embed error: {e}")
            await self.send_error(connection_id, f"Embedding failed: {str(e)}", request_id)

    async def handle_search(self, connection_id: str, data: Dict[str, Any], request_id: str):
        """Handle vector search request."""
        try:
            table_name = data.get("table_name", "default")
            query = data.get("query", "")
            limit = data.get("limit", 10)
            threshold = data.get("threshold", 0.0)
            provider = data.get("provider")

            if not query:
                await self.send_error(connection_id, "No query provided", request_id)
                return

            # Generate query embedding
            query_embedding = await embed_text(query, provider)

            # Perform search
            db_manager = await get_db_manager()
            results = await db_manager.search(table_name, query_embedding, limit, threshold)

            # Format results
            formatted_results = []
            for result in results:
                formatted_results.append({
                    "text": result.get("text", ""),
                    "metadata": json.loads(result.get("metadata", "{}")),
                    "score": result.get("_distance", 0.0),
                    "created_at": result.get("created_at", "").isoformat() if result.get("created_at") else None
                })

            response = {
                "type": "search_response",
                "request_id": request_id,
                "data": {
                    "results": formatted_results,
                    "count": len(formatted_results),
                    "query": query,
                    "table_name": table_name
                },
                "timestamp": datetime.now().isoformat()
            }

            await connection_manager.send_message(connection_id, response)

        except Exception as e:
            logger.error(f"WebSocket search error: {e}")
            await self.send_error(connection_id, f"Search failed: {str(e)}", request_id)

    async def handle_store(self, connection_id: str, data: Dict[str, Any], request_id: str):
        """Handle vector store request."""
        try:
            table_name = data.get("table_name", "default")
            texts = data.get("texts", [])
            vectors = data.get("vectors")
            metadata = data.get("metadata", [])
            provider = data.get("provider")

            if not texts:
                await self.send_error(connection_id, "No texts provided", request_id)
                return

            # Generate embeddings if vectors not provided
            if vectors is None:
                embeddings = await embed_texts(texts, provider)
            else:
                embeddings = vectors

            # Prepare data for insertion
            insert_data = []
            for i, (text, embedding) in enumerate(zip(texts, embeddings)):
                item = {
                    "vector": embedding,
                    "text": text,
                    "metadata": json.dumps(metadata[i]) if metadata and i < len(metadata) else "{}",
                    "created_at": datetime.now()
                }
                insert_data.append(item)

            # Insert data
            db_manager = await get_db_manager()
            success = await db_manager.insert_data(table_name, insert_data)

            response = {
                "type": "store_response",
                "request_id": request_id,
                "data": {
                    "success": success,
                    "stored_count": len(insert_data) if success else 0,
                    "table_name": table_name
                },
                "timestamp": datetime.now().isoformat()
            }

            await connection_manager.send_message(connection_id, response)

        except Exception as e:
            logger.error(f"WebSocket store error: {e}")
            await self.send_error(connection_id, f"Store failed: {str(e)}", request_id)

    async def handle_health(self, connection_id: str, data: Dict[str, Any], request_id: str):
        """Handle health check request."""
        try:
            # Check database
            db_manager = await get_db_manager()
            db_stats = await db_manager.get_database_stats()
            db_healthy = "error" not in db_stats

            # Check embedding service
            embedding_service = await get_embedding_service()
            provider_health = await embedding_service.health_check()
            embedding_healthy = any(provider_health.values())

            response = {
                "type": "health_response",
                "request_id": request_id,
                "data": {
                    "status": "healthy" if db_healthy and embedding_healthy else "degraded",
                    "database": {
                        "healthy": db_healthy,
                        "stats": db_stats if db_healthy else {}
                    },
                    "embedding": {
                        "healthy": embedding_healthy,
                        "providers": provider_health
                    }
                },
                "timestamp": datetime.now().isoformat()
            }

            await connection_manager.send_message(connection_id, response)

        except Exception as e:
            logger.error(f"WebSocket health check error: {e}")
            await self.send_error(connection_id, f"Health check failed: {str(e)}", request_id)

    async def handle_stats(self, connection_id: str, data: Dict[str, Any], request_id: str):
        """Handle stats request."""
        try:
            stats = connection_manager.get_connection_stats()

            response = {
                "type": "stats_response",
                "request_id": request_id,
                "data": stats,
                "timestamp": datetime.now().isoformat()
            }

            await connection_manager.send_message(connection_id, response)

        except Exception as e:
            logger.error(f"WebSocket stats error: {e}")
            await self.send_error(connection_id, f"Stats failed: {str(e)}", request_id)

    async def handle_ping(self, connection_id: str, data: Dict[str, Any], request_id: str):
        """Handle ping request."""
        response = {
            "type": "pong",
            "request_id": request_id,
            "data": {"timestamp": datetime.now().isoformat()},
            "timestamp": datetime.now().isoformat()
        }
        await connection_manager.send_message(connection_id, response)

    async def handle_subscribe(self, connection_id: str, data: Dict[str, Any], request_id: str):
        """Handle subscription request."""
        try:
            subscription_type = data.get("type")
            if not subscription_type:
                await self.send_error(connection_id, "No subscription type provided", request_id)
                return

            if connection_id not in self.subscriptions:
                self.subscriptions[connection_id] = set()

            self.subscriptions[connection_id].add(subscription_type)

            response = {
                "type": "subscribe_response",
                "request_id": request_id,
                "data": {
                    "subscription_type": subscription_type,
                    "active_subscriptions": list(self.subscriptions[connection_id])
                },
                "timestamp": datetime.now().isoformat()
            }

            await connection_manager.send_message(connection_id, response)

        except Exception as e:
            logger.error(f"WebSocket subscribe error: {e}")
            await self.send_error(connection_id, f"Subscribe failed: {str(e)}", request_id)

    async def handle_unsubscribe(self, connection_id: str, data: Dict[str, Any], request_id: str):
        """Handle unsubscription request."""
        try:
            subscription_type = data.get("type")
            if not subscription_type:
                await self.send_error(connection_id, "No subscription type provided", request_id)
                return

            if connection_id in self.subscriptions:
                self.subscriptions[connection_id].discard(subscription_type)

                response = {
                    "type": "unsubscribe_response",
                    "request_id": request_id,
                    "data": {
                        "subscription_type": subscription_type,
                        "active_subscriptions": list(self.subscriptions[connection_id])
                    },
                    "timestamp": datetime.now().isoformat()
                }

                await connection_manager.send_message(connection_id, response)

        except Exception as e:
            logger.error(f"WebSocket unsubscribe error: {e}")
            await self.send_error(connection_id, f"Unsubscribe failed: {str(e)}", request_id)

    async def send_error(self, connection_id: str, error_message: str, request_id: str = None):
        """Send error message to connection."""
        error_response = {
            "type": "error",
            "request_id": request_id,
            "data": {
                "error": error_message,
                "timestamp": datetime.now().isoformat()
            },
            "timestamp": datetime.now().isoformat()
        }
        await connection_manager.send_message(connection_id, error_response)

    async def broadcast_event(self, event_type: str, data: Dict[str, Any]):
        """Broadcast event to subscribed connections."""
        message = {
            "type": f"event_{event_type}",
            "data": data,
            "timestamp": datetime.now().isoformat()
        }

        # Find subscribed connections
        subscribed_connections = []
        for connection_id, subscriptions in self.subscriptions.items():
            if event_type in subscriptions:
                subscribed_connections.append(connection_id)

        if subscribed_connections:
            await connection_manager.broadcast_message(message, exclude=set())


# Global WebSocket handler
websocket_handler = WebSocketHandler()


async def websocket_endpoint(websocket: WebSocket):
    """WebSocket endpoint for real-time vector operations."""
    connection_id = None

    try:
        # Accept connection
        connection_id = await connection_manager.connect(websocket)

        # Send welcome message
        welcome_message = {
            "type": "welcome",
            "data": {
                "connection_id": connection_id,
                "service": "vectordb-microservice",
                "version": "1.0.0",
                "timestamp": datetime.now().isoformat()
            },
            "timestamp": datetime.now().isoformat()
        }
        await websocket.send_json(welcome_message)

        # Handle messages
        while True:
            try:
                # Receive message
                data = await websocket.receive_json()
                message = WebSocketMessage(**data)

                # Handle message
                await websocket_handler.handle_message(connection_id, message)

            except ValueError as e:
                # Invalid message format
                await websocket_handler.send_error(connection_id, f"Invalid message format: {e}")
            except Exception as e:
                logger.error(f"Error processing WebSocket message: {e}")
                await websocket_handler.send_error(connection_id, f"Message processing error: {e}")

    except WebSocketDisconnect:
        logger.info(f"WebSocket client disconnected: {connection_id}")
    except Exception as e:
        logger.error(f"WebSocket error: {e}")
    finally:
        if connection_id:
            connection_manager.disconnect(connection_id)


# Convenience functions for broadcasting events
async def broadcast_new_vector(table_name: str, count: int):
    """Broadcast new vector storage event."""
    await websocket_handler.broadcast_event("new_vectors", {
        "table_name": table_name,
        "count": count,
        "timestamp": datetime.now().isoformat()
    })


async def broadcast_search_completed(table_name: str, query: str, result_count: int):
    """Broadcast search completion event."""
    await websocket_handler.broadcast_event("search_completed", {
        "table_name": table_name,
        "query": query,
        "result_count": result_count,
        "timestamp": datetime.now().isoformat()
    })


async def broadcast_table_created(table_name: str):
    """Broadcast table creation event."""
    await websocket_handler.broadcast_event("table_created", {
        "table_name": table_name,
        "timestamp": datetime.now().isoformat()
    })


async def broadcast_table_deleted(table_name: str):
    """Broadcast table deletion event."""
    await websocket_handler.broadcast_event("table_deleted", {
        "table_name": table_name,
        "timestamp": datetime.now().isoformat()
    })